from pathlib import Path
import numpy as np
import cv2
import mss
import monitor
from logger import logger as log

#Generic function to see if something is present on the screen
def find(path, confidence=0.9, return_cords=False, center_on_found=True):
    """
    在给定路径中查找图像，可选地返回其坐标或仅返回是否找到的布尔值。
    
    参数:
    - path: 图像文件的路径。
    - confidence: 识别图像所需的最低置信度（0到1之间）。
    - return_cords: 是否返回图像的坐标。如果为True，则返回坐标；如果为False，则返回一个布尔值，指示图像是否被找到。
    - center_on_found: 当return_cords为True时，如果找到图像，是否返回图像中心的坐标而不是左上角的坐标。
    
    返回值:
    - 如果return_cords为True，并且找到了图像，返回图像中心的坐标（如果center_on_found为True）或图像的左上角坐标及其宽度和高度。
    - 如果return_cords为False，返回一个布尔值，指示是否找到了图像。
    - 如果没有找到图像，且return_cords为True，返回None。
    
    异常:
    - 抛出任何在查找过程中遇到的异常。
    """

    try:
        if return_cords:
            cords = locate(path, confidence=confidence)  # 尝试定位图像并获取其坐标
            
            log.debug("Found cords: " + str(cords))  # 记录找到的坐标

            if cords is not None:
                left, top, width, height = cords
                if center_on_found:
                    return (left + width // 2, top + height // 2)  # 返回图像中心的坐标
                else:
                    return (left, top, width, height)  # 返回图像的左上角坐标及其尺寸
            else:
                return None  # 如果没有找到图像，返回None
        return True if locate(path, confidence=confidence) is not None else False  # 返回一个布尔值，指示是否找到了图像

    except Exception as e:
        raise Exception(e)  # 将查找过程中遇到的任何异常向上抛出


def load_image(img):
    """
    加载图像。
    
    根据输入参数的类型加载或转换图像，支持从文件路径、OpenCV numpy数组或PIL图像加载。
    如果输入是路径，则尝试以灰度模式读取图像；
    如果输入是numpy数组，则将其转换为灰度；
    如果输入是PIL图像，则先转换为RGB，然后转换为灰度。
    注意：不支持带alpha通道的图像，会自动转换为RGB或灰度。
    
    参数:
    img - 图像的路径（Path类型）、OpenCV numpy数组或PIL图像对象。
    
    返回值:
    img_cv - 加载或转换后的灰度图像（OpenCV numpy数组）。
    
    抛出异常:
    TypeError - 如果输入类型不是预期的之一。
    IOError - 如果无法读取图像文件。
    """
    # 根据图像来源进行不同的处理
    if isinstance(img, Path):
        # 从文件路径读取图像，若失败则抛出IOError
        img_cv = cv2.imread(str(img), cv2.IMREAD_GRAYSCALE)
        if img_cv is None:
            raise IOError(f"Failed to read {img} because file is missing, has improper permissions, or is an unsupported or invalid format")
    elif isinstance(img, np.ndarray):
        # 如果图像是numpy数组，转换为灰度
        img_cv = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    elif hasattr(img, 'convert'):
        # 如果图像是PIL图像，转换为OpenCV格式的numpy数组再转换为灰度
        img_array = np.array(img.convert('RGB'))
        img_cv = img_array[:, :, ::-1].copy()  # 将RGB转换为BGR
        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
    else:
        # 如果输入类型不符，抛出TypeError
        raise TypeError('expected an image filename, OpenCV numpy array, or PIL image')
    
    return img_cv


def locate_all(template_path, confidence=0.9, limit=100, region=None):
    """
    使用模板匹配方法，在通过mss库截取的屏幕快照中寻找模板图像的所有匹配位置。

    @template_path - 模板图像的路径
    @confidence - 匹配置信度阈值，取值范围为(0.0, 1.0)，默认为0.9
    @limit - 返回匹配结果的最大数量，默认为100
    @region - 搜索的区域，默认为None，即搜索整个屏幕

    返回值：一个包含所有匹配位置坐标的列表，如果未找到匹配项则返回None。

    说明：此函数依赖于pyscreeze和opencv库。通过调整模板图像的大小以适应屏幕分辨率，并使用cv2.matchTemplate方法寻找匹配项。
    """

    # 确定截取屏幕的区域
    screenshotArea = {
        'top': 0, 
        'left': 0, 
        'width': monitor.width, 
        'height': monitor.height
    } if region is None else region

    # 验证置信度参数是否在合理范围内
    if 0.0 > confidence <= 1.0:
        raise ValueError("Confidence must be a value between 0.0 and 1.0")

    with mss.mss() as screenshotter:
        # 将截取的屏幕区域转换为opencv图像对象
        img = np.array(screenshotter.grab(screenshotArea))
        screenshot = load_image(img) 

        # 如果指定了搜索区域，则调整截图范围
        if region:
            screenshot = screenshot[region[1]:region[1]+region[3],
                                    region[0]:region[0]+region[2]
                                    ]
        else:
            region = (0, 0)

        # 加载模板图像
        template = load_image(template_path)

        confidence = float(confidence)

        # 获取模板图像的宽和高
        templateHeight, templateWidth = template.shape[:2]

        # 根据屏幕分辨率调整模板图像的大小
        if monitor.width != 2560 or monitor.height != 1440:
            template = cv2.resize(
                template, 
                dsize=(int(templateWidth/(2560/monitor.width)), int(templateHeight/(1440/monitor.height))), 
                interpolation=cv2.INTER_AREA 
            )
        
        # 使用模板匹配查找所有匹配项
        result = cv2.matchTemplate(screenshot, template, cv2.TM_CCOEFF_NORMED)    # 生成模板和屏幕快照的热力图
        match_indices = np.arange(result.size)[(result > confidence).flatten()]
        matches = np.unravel_index(match_indices[:limit], result.shape)
        
        # 计算匹配区域的坐标
        matchesX = matches[1] * 1 + region[0]
        matchesY = matches[0] * 1 + region[1]

        # 如果未找到任何匹配项，则返回None
        if len(matches[0]) == 0:
            return None
        else:
            return [ (x, y, templateWidth, templateHeight) for x, y in zip(matchesX, matchesY) ]

def locate(template_path, confidence=0.9, tries=1):
    """
    在屏幕上定位一个模板。

    注意：当前@tries参数暂不生效

    参数:
    - template_path: 模板文件的路径，用于在屏幕上进行匹配。
    - confidence: 匹配的置信度，默认为0.9，即只有当匹配度高于此值时才认为是成功找到模板。
    - tries: 尝试匹配的次数，默认为1。注意，当前此参数没有实际作用。

    返回值:
    - 如果成功找到模板，则返回匹配结果的第一个元素；如果未找到或发生错误，则返回None。
    """
    # 尝试定位所有匹配的模板
    result = locate_all(template_path, confidence)
    # 如果有找到匹配的模板，返回第一个匹配结果，否则返回None
    return result[0] if result is not None else None